#include <cstdint>
#include <cassert>
#include <algorithm>
#include "Enigma.hpp"

#define N 16    // charset_size
#define M 3     // rotors_number


/// @brief Test whether those shifts are less than charset_size.
/// @param seed 
static void test_shifts(const std::uint64_t seed)
{
    // Given
    Enigma::Key<M> key = Enigma::GenerateKey<N, M>(seed);
    // Then
    assert(key.rotor_shifts.size() == M);
    assert(
        std::all_of(key.rotor_shifts.cbegin(), key.rotor_shifts.cend(),
            [](const std::uint8_t shift){ return shift < N; }
        )
    );
}


/// @brief Test whether those connect pairs are valid.
/// @param seed 
static void test_connect_pairs(const std::uint64_t seed)
{
    // Given
    Enigma::Key<M> key = Enigma::GenerateKey<N, M>(seed);
    // Then
    assert(key.connect_pair_list.size() == N / 2);

    // When
    std::array<std::uint8_t, N> mapping;
    std::iota(mapping.begin(), mapping.end(), 0);
    for (auto& pair : key.connect_pair_list) {
        // Then assert has no multi-mapping.
        assert(
            mapping[pair.first] == pair.first &&
            mapping[pair.second] == pair.second
        );

        mapping[pair.first] = pair.second;
        mapping[pair.second] = pair.first;
    }
}


/// @brief Test whether two keys with different seeds are different.
static void test_different_keys()
{
    // Given
    Enigma::Key<M> key_1 = Enigma::GenerateKey<N, M>(123);
    Enigma::Key<M> key_2 = Enigma::GenerateKey<N, M>(456);
    // Then
    assert(
        !std::equal(
            key_1.connect_pair_list.begin(),
            key_1.connect_pair_list.end(),
            key_2.connect_pair_list.begin()
        )
    );
}


/// @brief Test whether two keys with same seeds are same.
static void test_same_keys()
{
    // Given
    Enigma::Key<M> key_1 = Enigma::GenerateKey<N, M>(123);
    Enigma::Key<M> key_2 = Enigma::GenerateKey<N, M>(123);
    // Then
    assert(
        std::equal(
            key_1.connect_pair_list.begin(),
            key_1.connect_pair_list.end(),
            key_2.connect_pair_list.begin()
        )
    );
}


int main()
{
    test_shifts(123);
    test_shifts(456);
    test_shifts(789);

    test_connect_pairs(123);
    test_connect_pairs(456);
    test_connect_pairs(789);


    test_different_keys();
    test_same_keys();
    return 0;
}